# encoding: UTF-8

import platform
import time as time

import numpy as np
import pandas as pd

import indicator.zpd_talib as zw

from const_struct.constant import *

MATYPE_MA = 'MA'  # 常规均线
MATYPE_EMA = 'EMA'  # 指数滑动均线
MATYPE_AMA = 'AMA'  # 自适应均线
"""
包含常用指标

"""


class TrendData(object):
    """趋势数据"""

    def __init__(self):
        self.datetime = None  # 格式：%Y%m%d %H:%M:%S
        self.ubl = TREND_POSITION_NONE  # 走势标识
        self.uprice = EMPTY_FLOAT  # 状态开始时的高价
        self.lprice = EMPTY_FLOAT  # 状态开始时的低价
        self.uprice1 = EMPTY_FLOAT  # 状态持续时的最高价
        self.lprice1 = EMPTY_FLOAT  # 状态持续时的最低价
        self.uperiod = EMPTY_INT  # 涨状态开始到当前的周期数
        self.lperiod = EMPTY_INT  # 跌状态开始到当前的周期数
        self.uperiod1 = EMPTY_INT  # 最高价时点到当前的周期数
        self.lperiod1 = EMPTY_INT  # 最低价时点到当前的周期数
        self.cprice = EMPTY_FLOAT  # 中间价
        self.hGTCount = EMPTY_INT  # 最高价比 uprice 大的频数
        self.lLTCount = EMPTY_INT  # 最低价比 lprice 小的频数
        self.bot = EMPTY_FLOAT  # 根据均衡状态的要求计算出的底
        self.top = EMPTY_FLOAT  # 根据均衡状态的要求计算出的顶
        self.dif = EMPTY_FLOAT  # 当前最高最低价差
        self.adif = EMPTY_FLOAT
        self.stddif = EMPTY_FLOAT  # dif的标准差
        self.isUp = 0  #
        self.upOk = 0  #
        self.isDown = 0  #
        self.downOk = 0  #
        self.isBalance = 0
        self.balanceOk = 0
        self.difu = 0.0
        self.difl = 0.0
        self.difc = 0.0
        self.difu2 = 0.0
        self.difl2 = 0.0
        self.difc2 = 0.0
        self.preOKData = None  # 前一涨、跌、盘确认时的数据(datetime,close,ubl,high,low）


class MAIndicatorFrame():
    """
    MA指标类
    """

    def __init__(self, symbol, data=None, params=None):

        self.value = {'MA1': [], 'MA2': [], 'MA3': []}


class TrendIndicatorFrame():
    """
    趋势指标
    """

    def __init__(self, symbol, period=5):

        self.period = period

    def init_value(self):
        self.ubl = TREND_POSITION_NONE  # 走势标识
        self.uprice = EMPTY_FLOAT  # 状态开始时的高价
        self.lprice = EMPTY_FLOAT  # 状态开始时的低价
        self.uprice1 = EMPTY_FLOAT  # 状态持续时的最高价
        self.lprice1 = EMPTY_FLOAT  # 状态持续时的最低价
        self.uperiod = EMPTY_INT  # 状态开始到当前的周期数
        self.lperiod = EMPTY_INT
        self.uperiod1 = EMPTY_INT  # 最高价时点到当前的周期数
        self.lperiod1 = EMPTY_INT  # 最低价时点到当前的周期数
        self.cprice = EMPTY_FLOAT  # 中间价
        self.hGTCount = EMPTY_INT  # 最高价比 uprice 大的频数
        self.lLTCount = EMPTY_INT  # 最低价比 lprice 小的频数
        self.bot = EMPTY_FLOAT  # 根据均衡状态的要求计算出的底
        self.top = EMPTY_FLOAT  # 根据均衡状态的要求计算出的顶


def MA(data, period=5):
    return zw.MA(data, period)


def hhv(data, periods, center=0):
    """
    取得 data 在 periods 周期内的最大值，当 periods 为偶数且 center==0 时，前边取一半（不含当前位置），
    :param data:pd.Series
    :param periods:
    :param center: 0==center, 1==before, -1=after
    :return:
    """
    if center == 0:
        hh = data.rolling(window=periods, min_periods=1, center=True).max()
    elif center == 1:
        hh = data.rolling(window=periods, min_periods=1, center=False).max()
    else:  # center==-1:
        hh = data[::-1].rolling(window=periods, min_periods=1, center=False).max()[::-1]
    return hh


def llv(data, periods, center=0):
    """
    取得 data 在 periods 周期内的最小值，当 periods 为偶数且 center==0 时，前边取一半（不含当前位置），
    :param data:pd.Series
    :param periods:
    :param center: 0==center, 1==before, -1=after
    :return:
    """
    if center == 0:
        ll = data.rolling(window=periods, min_periods=1, center=True).min()
    elif center == 1:
        ll = data.rolling(window=periods, min_periods=1, center=False).min()
    else:  # center==-1:
        ll = data[::-1].rolling(window=periods, min_periods=1, center=False).min()[::-1]
    return ll


def fractal(data, high_sgn='high', low_sgn='low', periods=2):
    """
    计算上下分形
    :param data:pd.DataFrame,应该是 KBar 序列
    :param high_sgn: 要标记的高点列名
    :param low_sgn: 要标记的低点列名，可与 high_sgn 相同
    :param periods: 2*periods+1= 分形用周期
    :return:分形序列
    """
    pdf = 2 * periods + 1
    hh = hhv(data=data[high_sgn], periods=pdf)
    ll = llv(data=data[low_sgn], periods=pdf)


def NR(df, n=5, ksgn='hl'):
    """
    calculate the n-periods range according to ksgn
    :param df: DataFrame,columns=['datetime','open','high','low','close','vol']
    :param n: periods
    :param ksgn: 'hl' or 'lh' calculate max(high,n)-min(low,n)
    :return:
    """
    if ksgn.lower() == 'hl' or ksgn.lower() == 'lh':
        high = df['high']
        low = df['low']
    else:
        high = df[ksgn]
        low = df[ksgn]
    hh = hhv(high, n, center=1)
    ll = llv(low, n, center=1)
    tr = hh - ll
    return tr


def Speed(data, spp):
    """速度：
    """
    ret = []
    if data is None or spp <= 0:
        return None
    ret.append(0.0)
    for i in range(1, min(spp, len(data))):
        ret.append(data[i] - data[0])
    for i in range(spp, len(data)):
        ret.append(data[i] - data[i - spp])
    return ret


def TR(df, sgnClose='close', sgnHigh='high', sgnLow='low', n=5):
    """
    计算真实波幅
    【输入】
        df, pd.dataframe格式数据源
        sgnClose，列名，一般是：close收盘价
    【输出】
        pd.dataframe格式数据源,

    """

    i = 0
    TR_l = [df[sgnHigh][i] - df[sgnLow][0]]
    while i < len(df) - 1:  # df.index[-1]:
        TR = max(df[sgnHigh].iloc[i + 1], df[sgnClose].iloc[i]) - min(df[sgnLow].iloc[i + 1], df[sgnClose].iloc[i])
        TR_l.append(TR)
        i += 1
    TR_s = pd.Series(TR_l, index=df.index, name='TR')
    ATR = pd.Series(TR_s.ewm(span=n, min_periods=n).mean(), name='ATR', index=df.index)
    return TR_s, ATR


def ACSP(dt, data, spp=5, acp=5):
    """速度加速度
    速度：价格的离差
    加速度：速度的离差"""


def TrendClass(df, timeWindow=5, sgnDatetime='datetime', sgnOpen='open', sgnHigh='high',
               sgnLow='low', sgnClose='close', sgnVol='volume', rangePCT=0, riskTimes=1.5, dif_times=4.0):
    """
    走势分类：将一段K线走势分为涨、盘、跌，也可称为趋势状态（涨、跌）和均衡状态。
    均衡状态是指在给定的时间窗口T和最大波动幅度R，HHV(HIGH,T)－LLV(LOW,T) <CC * R,CC=(HHV(HIGH,T)+LLV(LOW,T))/2。
    :param df: dataFrame,dt='datetime',open='open',high='high',low='low',close='close',vol='volume'是df的结构
    :param timeWindow: 均衡状态形成的最小时间窗口，原则上不小于3
    :param rangePCT:给定的均衡区间振荡幅度（%%，为整数），若为0，则由序列本身确定
    :param riskTimes:风险乘数，当rangePCT==0时，均衡区间的范围由用到的风险单位*该风险乘数，风险通常用真实波幅
    :return:TrendData 的序列

    一段K线，K={Kt，t=0,1,...,N},可分为均衡状态、上涨状态、下跌状态，每种状态应该由4类K线组成：开始、确认、持续和结束，
    也可能没有持续。


    """
    if timeWindow < 3:
        return None

    op = df[sgnOpen]
    hp = df[sgnHigh]
    lp = df[sgnLow]
    cp = df[sgnClose]
    # volume=df[sgnVol]
    dts = df.index[0]  # [sgnDatetime]

    rtn = []
    td = TrendData()

    i = 0
    o = op[i]
    h = hp[i]
    l = lp[i]
    c = cp[i]
    d = dts
    cc = (h + l) / 2.0
    dif = h - l
    adif = dif
    hh = cc + dif
    ll = cc - dif
    difhc = h - c
    difcl = c - l
    td.datetime = dts
    td.lperiod = 0
    td.lperiod1 = 0
    td.lprice = l
    td.lprice1 = l
    td.ubl = TREND_POSITION_BALANCE_START
    td.uperiod = 0
    td.uperiod1 = 0
    td.uprice = h
    td.uprice1 = h
    td.cprice = cc
    td.hGTCount = 0
    td.lLTCount = 0
    td.bot = cc - adif * dif_times  # (1 + rangePCT / 10000.0) / 2.0
    td.top = cc + adif * dif_times  # (1 + rangePCT / 10000.0) / 2.0
    td.dif = dif
    td.adif = adif
    td.stddif = 0

    if rangePCT == 0 or (dif < cc * rangePCT):
        td.ubl = TREND_POSITION_BALANCE_START
        #         td.balanceOk = 1
        td.isBalance = 1
    else:
        if difhc > difcl:  # 上影大于下影，定义为下跌
            td.ubl = TREND_POSITION_DOWN_START
            #             td.downOk = 1
            td.isDown = 1
        else:  # 上影小于下影，定义为上涨
            td.ubl = TREND_POSITION_UP_START
            #             td.upOk = 1
            td.isUp = 1
    td.preOKData = {'datetime': d, 'ubl': td.ubl, 'close': c, 'high': h, 'low': l}
    rtn.append(td)

    while i < len(df) - 1:
        i += 1
        d = df.index[i]
        newK = {'datetime': d, 'high': hp[i], 'low': lp[i], 'close': cp[i]}
        td = newTrendClass(rtn[i - 1], newK, timeWindow=timeWindow,
                           dif_times=dif_times, rangePCT=rangePCT, riskTimes=riskTimes)
        rtn.append(td)
    return rtn


def newTrendClass(td0, newK, timeWindow=15, rangePCT=0, riskTimes=1.5, dif_times=4.0):
    '''
    根据新的K值，计算新的趋势数据
    :param td0: 原来的趋势数据，类型TrendData
    :param newK: 新的k值，包含datetime,high,low,close 的字典
    :param rangePCT: 给定的均衡区间振荡幅度（%%，为整数），若为0，则由序列本身确定
    :param dif_times:
    :return: 类型TrendData
    '''
    h = newK['high']
    l = newK['low']
    c = newK['close']
    d = newK['datetime']
    ubl = td0.ubl
    cc0 = td0.cprice
    dif0 = td0.dif
    l0 = td0.lprice1
    h0 = td0.uprice1
    top0 = td0.top
    bot0 = td0.bot
    stddif0 = td0.stddif
    dif = max(cc0, h) - min(cc0, l)
    adif0 = td0.adif
    adif = (adif0 + dif) / 2.0
    stddif = np.sqrt(np.square(dif - adif) / 2.0 + np.square(dif0 - adif) / 2.0)

    hh = max(h, td0.uprice1)
    ll = min(l, td0.lprice1)
    cc = (hh + ll) / 2.0
    # dif = hh - ll
    top = max(cc0 + adif0 / 2 + riskTimes * stddif0, top0)
    bot = min(cc0 - adif0 / 2 - riskTimes * stddif0, bot0)

    td = TrendData()
    td.datetime = d
    td.lperiod = td0.lperiod + 1
    td.lperiod1 = td0.lperiod1 + 1
    td.lprice = td0.lprice
    td.lprice1 = td0.lprice1
    td.ubl = td0.ubl
    td.uperiod = td0.uperiod + 1
    td.uperiod1 = td0.uperiod1 + 1
    td.uprice = td0.uprice
    td.uprice1 = td0.uprice1
    td.cprice = (cc + cc0) / 2.0
    td.top = top
    td.bot = bot
    td.dif = dif
    td.adif = adif
    td.stddif = stddif

    statusMsg = None

    if ubl > TREND_POSITION_UP_END:  # 前期是均衡状态
        if h > td.uprice1:  # 创新高
            td.uprice1 = h
            td.uperiod1 = 0
            td.hGTCount += 1
        if l < td.lprice1:  # 创新低
            td.lprice1 = l
            td.lperiod1 = 0
            td.lLTCount += 1

        if (h < top and l < bot) or \
                (h > top and l < bot and \
                 h - c > c - l):
            # 均衡状态结束，进入下跌状态
            statusMsg = (d, '均衡状态结束，进入下跌状态')
            td.ubl = TREND_POSITION_DOWN_START
            td.uprice = h
            td.uprice1 = h
            td.lprice = l
            td.lprice1 = l
            td.lperiod = 0
            td.lperiod1 = 0
            td.uperiod = 0
            td.uperiod1 = 0
            cc = (h + l) / 2.0
            td.cprice = cc
            td.hGTCount = 0
            td.lLTCount = 0
            td.dif = h - l
            td.adif = td.dif
            td.stddif = 0
            td.top = cc + td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.bot = cc - td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.isDown = 1

            if td0.preOKData['close'] < c:
                td.downOk = 1
            else:
                td.upOk = 1

            td.difu = td.uprice - td0.uprice
            td.difl = td.lprice - td0.lprice
            td.difc = td.cprice - td0.cprice
            td.difu2 = td.difu - td0.difu
            td.difl2 = td.difl - td0.difl
            td.difc2 = td.difc - td0.difc
            td.preOKData = {'datetime': d, 'ubl': td.ubl, 'close': c, 'high': h, 'low': l}

        elif (h > top and l >= bot) or \
                (h > top and l < bot and \
                 h - c < c - l):  # 均衡状态结束，进入上涨状态
            statusMsg = (d, '均衡状态结束，进入上涨状态')
            td.ubl = TREND_POSITION_UP_START
            td.uprice = h
            td.uprice1 = h
            td.lprice = l
            td.lprice1 = l
            td.lperiod = 0
            td.lperiod1 = 0
            td.uperiod = 0
            td.uperiod1 = 0
            cc = (h + l) / 2.0
            td.cprice = cc
            td.hGTCount = 0
            td.lLTCount = 0
            td.dif = h - l
            td.adif = td.dif
            td.stddif = 0
            td.top = cc + td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.bot = cc - td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.isUp = 1

            if td0.preOKData['close'] > c:
                td.downOk = 1
            else:
                td.upOk = 1

            td.difu = td.uprice - td0.uprice
            td.difl = td.lprice - td0.lprice
            td.difc = td.cprice - td0.cprice
            td.difu2 = td.difu - td0.difu
            td.difl2 = td.difl - td0.difl
            td.difc2 = td.difc - td0.difc
            td.preOKData = {'datetime': d, 'ubl': td.ubl, 'close': c, 'high': h, 'low': l}

        else:  # 均衡状态持续
            if td0.ubl == TREND_POSITION_BALANCE_START:
                td.ubl = TREND_POSITION_BALANCE + TREND_POSITION_BALANCE_START  # 默认为均衡状态持续
            elif td0.ubl == TREND_POSITION_BALANCE_OK:
                td.ubl = TREND_POSITION_BALANCE + TREND_POSITION_BALANCE_OK
            else:
                td.ubl = td0.ubl
                # 均衡状态确认后，直到均衡状态被破坏，上、中、下线不再改变
            if (td0.ubl & TREND_POSITION_BALANCE_OK) == TREND_POSITION_BALANCE_OK:
                td.cprice = td0.cprice
                td.top = td0.top
                td.bot = td0.bot

            td.preOKData = td0.preOKData
            td.cprice = cc0
            if td.lperiod == timeWindow - 1 or td.uperiod == timeWindow - 1:
                td.ubl = TREND_POSITION_BALANCE_OK  # 均衡状态确认

    elif ubl > TREND_POSITION_DOWN_END:  # 前期是上涨状态
        if l < bot:  # 上涨结束，下跌开始
            statusMsg = (d, '上涨结束，下跌开始')
            td.ubl = TREND_POSITION_DOWN_START
            td.uprice = h
            td.uprice1 = h
            td.lprice = l
            td.lprice1 = l
            td.lperiod = 0
            td.lperiod1 = 0
            td.uperiod = 0
            td.uperiod1 = 0
            cc = (h + l) / 2.0
            td.cprice = cc
            td.hGTCount = 0
            td.lLTCount = 0
            td.dif = h - l
            td.adif = td.dif
            td.stddif = 0
            td.top = cc + td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.bot = cc - td.dif * dif_times  ##* (0.5 + rangePCT / 10000.0)
            td.isDown = 1

            if td0.preOKData['close'] > c:
                td.downOk = 1
            else:
                td.upOk = 1

            td.difu = td.uprice - td0.uprice
            td.difl = td.lprice - td0.lprice
            td.difc = td.cprice - td0.cprice
            td.difu2 = td.difu - td0.difu
            td.difl2 = td.difl - td0.difl
            td.difc2 = td.difc - td0.difc
            td.preOKData = {'datetime': d, 'ubl': td.ubl, 'close': c, 'high': h, 'low': l}

        elif h < td0.uprice1 and td0.uperiod1 == timeWindow - 1:  # period观察窗口内未创新高，新的均衡状态开始
            statusMsg = (d, '观察窗口内未创新高，新的均衡状态开始')
            td.ubl = TREND_POSITION_BALANCE_START
            td.uprice = h
            td.uprice1 = h
            td.lprice = l
            td.lprice1 = l
            td.lperiod = 0
            td.lperiod1 = 0
            td.uperiod = 0
            td.uperiod1 = 0
            cc = (h + l) / 2.0
            td.cprice = cc
            td.hGTCount = 0
            td.lLTCount = 0
            td.dif = h - l
            td.adif = td.dif
            td.stddif = 0
            td.top = cc + td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.bot = cc - td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.isBalance = 1
            td.balanceOk = 1

            if td0.preOKData['close'] > c:
                td.downOk = 1
            else:
                td.upOk = 1

            td.difu = td.uprice - td0.uprice
            td.difl = td.lprice - td0.lprice
            td.difc = td.cprice - td0.cprice
            td.difu2 = td.difu - td0.difu
            td.difl2 = td.difl - td0.difl
            td.difc2 = td.difc - td0.difc
            td.preOKData = {'datetime': d, 'ubl': td.ubl, 'close': c, 'high': h, 'low': l}

        else:  # 上涨持续
            td.ubl = TREND_POSITION_UP
            td.preOKData = td0.preOKData

            if h > h0:  # 创新高，
                # il = max(0, i - period)
                # mn = min(lp[il:i])
                td.lprice1 = l
                td.lperiod1 = 0
                td.uprice1 = h
                td.uperiod1 = 0
                td.hGTCount += 1
                cc = (h + l) / 2.0
                td.cprice = cc
                td.dif = h - l
                td.adif = td.dif
                td.stddif = 0
                td.top = cc + td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
                td.bot = cc - td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)

    else:  # 前期是下跌状态
        if h > top:  # 下跌结束，上涨开始
            statusMsg = (d, '下跌结束，上涨开始')
            td.ubl = TREND_POSITION_UP_START
            td.uprice = h
            td.uprice1 = h
            td.lprice = l
            td.lprice1 = l
            td.lperiod = 0
            td.lperiod1 = 0
            td.uperiod = 0
            td.uperiod1 = 0
            cc = (h + l) / 2.0
            td.cprice = cc
            td.hGTCount = 0
            td.lLTCount = 0
            td.dif = h - l
            td.adif = td.dif
            td.stddif = 0
            td.top = cc + td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.bot = cc - td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.isUp = 1

            if td0.preOKData['close'] > c:
                td.downOk = 1
            else:
                td.upOk = 1

            td.difu = td.uprice - td0.uprice
            td.difl = td.lprice - td0.lprice
            td.difc = td.cprice - td0.cprice
            td.difu2 = td.difu - td0.difu
            td.difl2 = td.difl - td0.difl
            td.difc2 = td.difc - td0.difc
            td.preOKData = {'datetime': d, 'ubl': td.ubl, 'close': c, 'high': h, 'low': l}


        elif l > td0.lprice1 and td0.lperiod1 == timeWindow - 1:  # period观察窗口内未创新低，新的均衡状态开始
            statusMsg = (d, '观察窗口内未创新低，新的均衡状态开始')
            td.ubl = TREND_POSITION_BALANCE_START
            td.uprice = h
            td.uprice1 = h
            td.lprice = l
            td.lprice1 = l
            td.lperiod = 0
            td.lperiod1 = 0
            td.uperiod = 0
            td.uperiod1 = 0
            cc = (h + l) / 2.0
            td.cprice = cc
            td.hGTCount = 0
            td.lLTCount = 0
            td.dif = h - l
            td.adif = td.dif
            td.stddif = 0
            td.top = cc + td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.bot = cc - td.dif * dif_times  # #* (0.5 + rangePCT / 10000.0)
            td.balanceOk = 1
            td.isBalance = 1

            if td0.preOKData['close'] > c:
                td.downOk = 1
            else:
                td.upOk = 1

            td.difu = td.uprice - td0.uprice
            td.difl = td.lprice - td0.lprice
            td.difc = td.cprice - td0.cprice
            td.difu2 = td.difu - td0.difu
            td.difl2 = td.difl - td0.difl
            td.difc2 = td.difc - td0.difc
            td.preOKData = {'datetime': d, 'ubl': td.ubl, 'close': c, 'high': h, 'low': l}

        else:  # 下跌持续
            td.ubl = TREND_POSITION_DOWN
            td.preOKData = td0.preOKData
            if l < l0:  # 创新低，
                # il=max(0,i-period)
                # mx = max(hp[il:i])
                td.uprice1 = h
                td.uperiod1 = 0
                td.lprice1 = l
                td.lperiod1 = 0
                td.lLTCount += 1
                cc = (h + l) / 2.0
                td.cprice = cc
                td.hGTCount = 0
                td.dif = h - l
                td.adif = td.dif
                td.stddif = 0
                td.top = cc + td.dif * dif_times  # * (0.5 + rangePCT / 10000.0)
                td.bot = cc - td.dif * dif_times  ##* (0.5 + rangePCT / 10000.0)

    # if statusMsg:print('newk=',newK,'td0=',td0.__dict__,adif,statusMsg)
    return td


def AMA(df, pm=10, pf=2, ps=30, ksgn='close'):
    """
    自适应均线
    :param df:pd.DataFrame, or pd.Series
    :param pm:
    :param pf:
    :param ps:
    :param ksgn: when df is DataFrame, ksgn is the column name on which to calculate
    :return:
    """
    if isinstance(df, pd.Series):
        price = df
    else:
        price = pd.Series(df[ksgn])
    dlen = len(price)
    P1 = pm
    DIR = price.diff(P1)
    VIR = abs(price.diff(1)).rolling(window=P1, min_periods=1).sum()
    for i in range(pm):
        DIR[i] = price[i] - price[0]
    VIR[0] = 0
    ER = pd.Series(DIR.div(VIR))
    ER[ER.isna()] = 0
    # print(DIR.head(), VIR.head(), ER.head())
    # print(len(ER),dlen)

    x = 2.0 / (pf + 1) - 2.0 / (ps + 1)
    CS = ER * x
    CS = CS.add(2.0 / (ps + 1))
    CQ = CS * CS
    i = 0
    adaptive = pd.Series([0] * dlen, index=df.index)
    while i < pm:
        adaptive[i] = price[0:i + 1].sum() / (i + 1)
        i = i + 1
    while i < dlen:
        adaptive[i] = adaptive[i - 1] * (1 - CQ[i]) + price[i] * (CQ[i])
        i += 1
    # print(CS,adaptive)
    # adaptive=EMA(DMA(TPRICE, CQ), 2)
    return pd.Series(adaptive)


def new_ama(pre_ama, price, pm=10, pf=2, ps=30):
    """

    :param pre_ama:
    :param price:
    :param pm:
    :param pf:
    :param ps:
    :return:
    """
    DIR = price[-1] - price[0 - pm]
    VIR = sum([abs(price[i + 1] - price[i]) for i in range(len(price) - 1)])
    ER = DIR / VIR
    CS = ER * (2 / (pf + 1) - 2 / (ps + 1)) + 2 / (ps + 1)
    alpha = CS ** 2
    return pre_ama * (1 - alpha) + price[-1] * alpha


def ZIG(df, n, ksgn='close'):
    """
    寻找满足条件的峰P、谷V,P V满足(P-V)*2/(P+V)>n
    :param df:pd.dataframe格式数据源,正常包含datetime,open,high,low,close,vol
    :param n:峰、谷之差的最低标准
    :param ksgn: 当为'hl' or 'lh' 时，峰取H，谷取L，其它取ksgn
    :return:pd.dataframe，包含：peak,peak_is,peak_ok,peak_pos,valley,valley_is,valley_ok,valley_pos,pv
    """
    data_len = len(df)
    p = [0] * data_len  # peak
    pp = [0] * data_len  # peak position
    pok = [0] * data_len  # 确认 peak
    pis = [0] * data_len  #
    v = [0] * data_len  # valley
    vp = [0] * data_len  # valley position
    vok = [0] * data_len  #
    vis = [0] * data_len  #
    pok_p = [0] * data_len  # 峰确认位置
    pok_v = [0] * data_len  # 峰确认价格
    vok_p = [0] * data_len  # 谷确认位置
    vok_v = [0] * data_len  # 谷确认价格

    curPV = 0  # 1=peak,-1=valley, 0=unknown

    if ksgn.upper() == 'HL' or ksgn.upper() == 'LH':
        h = df['high']
        l = df['low']
    else:
        h = df[ksgn]
        l = df[ksgn]

    i = 0
    p[i] = h[i]
    v[i] = l[i]
    pp[i] = i
    vp[i] = i
    while i < data_len - 1:
        p0 = p[i]
        pp0 = pp[i]
        v0 = v[i]
        vp0 = vp[i]
        i += 1
        p[i] = p[i - 1]
        pp[i] = pp[i - 1]
        v[i] = v[i - 1]
        vp[i] = vp[i - 1]
        if curPV == 0:
            if p0 < h[i]:  # 仓新高
                p[i] = h[i]
                pp[i] = i
                if (p[i] - v[i]) > v[i] * n:  # 谷被确认
                    vis[vp[i]] = 1
                    vok[i] = 1
                    v[i] = l[i]
                    vp[i] = i
                    curPV = -1
                    vok_p[i] = i
                    vok_v[i] = p[i]
            if v0 > l[i]:  # 仓新低
                v[i] = l[i]
                vp[i] = i
                if (p[i] - v[i]) > p[i] * n:  # 峰被确认
                    pis[pp[i]] = 1
                    pok[i] = 1
                    p[i] = h[i]
                    pp[i] = i
                    curPV = 1
                    pok_p[i] = i
                    pok_v[i] = v[i]
        elif curPV == 1:
            if p0 < h[i]:  # 仓新高
                p[i] = h[i]
                pp[i] = i
                if (p[i] - v[i]) > v[i] * n:  # 谷被确认
                    vis[vp[i]] = 1
                    vok[i] = 1
                    v[i] = l[i]
                    vp[i] = i
                    curPV = -1
                    vok_p[i] = i
                    vok_v[i] = p[i]
            if v0 > l[i]:  # 仓新低
                p[i] = h[i]
                pp[i] = i
                v[i] = l[i]
                vp[i] = i
        else:  # curPV==-1
            if p0 < h[i]:  # 仓新高
                p[i] = h[i]
                pp[i] = i
                v[i] = l[i]
                vp[i] = i
            if v0 > l[i]:  # 仓新低
                v[i] = l[i]
                vp[i] = i
                if (p[i] - v[i]) > p[i] * n:  # 峰被确认
                    pis[pp[i]] = 1
                    pok[i] = 1
                    p[i] = h[i]
                    pp[i] = i
                    curPV = 1
                    pok_p[i] = i
                    pok_v[i] = v[i]

    peak_is = pd.Series(pis, index=df.index, name='peak_is')
    peak_ok = pd.Series(pok, index=df.index, name='peak_ok')
    peak = pd.Series(p, index=df.index, name='peak')
    peak_pos = pd.Series(pp, index=df.index, name='peak_pos')

    valley_is = pd.Series(vis, index=df.index, name='valley_is')
    valley_ok = pd.Series(vok, index=df.index, name='valley_ok')
    valley = pd.Series(v, index=df.index, name='valley')
    valley_pos = pd.Series(vp, index=df.index, name='valley_pos')
    pok_v = pd.Series(pok_v, index=df.index, name='p_o_v')
    pok_p = pd.Series(pok_p, index=df.index, name='p_o_p')
    vok_v = pd.Series(vok_v, index=df.index, name='v_o_v')
    vok_p = pd.Series(vok_p, index=df.index, name='v_o_p')
    df[peak.name] = peak
    df[peak_is.name] = peak_is
    df[peak_ok.name] = peak_ok
    df[peak_pos.name] = peak_pos
    df[valley.name] = valley
    df[valley_is.name] = valley_is
    df[valley_ok.name] = valley_ok
    df[valley_pos.name] = valley_pos
    df['peak_ok_pos'] = pok_p
    df['peak_ok_value'] = pok_v
    df['valley_ok_pos'] = vok_p
    df['valley_ok_value'] = vok_v

    pv_is = peak_is - valley_is
    pv_ok = peak_ok - valley_ok
    pv = (peak + valley) / 2
    pv_ok_value = df['peak_ok_value'] + df['valley_ok_value']
    pv_ok_pos = df['peak_ok_pos']
    print(len(pok_v), len(vok_v), len(pv_ok_value))

    df['pv_ok_value'] = pv_ok_value
    df['pv_ok_pos'] = pv_ok_pos
    df['pv_ok'] = pv_ok
    df['pv'] = pd.Series(pv, index=df.index, name='pv')
    # df['pv_is']=df.query('valley_is==1')['valley_is']
    df['pv_is'] = pd.Series(pv_is, index=df.index, name='pv_is')
    df['pv_pos'] = peak_pos
    diff_peak = peak[peak_is == 1].diff(1)
    df['d_p'] = diff_peak
    df['d_p_p'] = peak_pos[peak_is == 1].diff(1)
    df['d_p_o_v'] = pok_v[peak_ok == 1].diff(1)
    df['d_p_o_p'] = pok_p[peak_ok == 1].diff(1)
    diff_valley = valley[valley_is == 1].diff(1)
    df['d_v'] = diff_valley
    df['d_v_p'] = valley_pos[valley_is == 1].diff(1)
    df['d_v_o_v'] = vok_v[valley_ok == 1].diff(1)
    df['d_v_o_p'] = vok_p[valley_ok == 1].diff(1)
    diff_pv = pv[pv_is != 0].diff(1)
    df['d_pv'] = diff_pv
    df['d_pv_p'] = df[pv_is != 0]['pv_pos'].diff(1)
    df['d_pv_o_v'] = pv_ok_value[pv_ok != 0].diff(1)
    df['d_pv_o_p'] = pv_ok_pos[pv_ok != 0].diff(1)
    return df


HHV = hhv
LLV = llv


def get_data():
    # HS300  ZZ500  万得全A 创业板指 上证综指 深圳成指
    if 'Windows' in platform.uname():
        path = r'd:\v1\notebook\ZZ500'
    else:
        path = r'../notebook/HS300'
    # os.chdir(path)
    fileName = path  # r'HS300.xlsx'
    # hs300 = createpickle(fileName)
    # fileName = r'/home/samba/secured/quantWeb/notebook/HS300'
    # print(fileName)
    data = pd.read_pickle(fileName)

    return data


def ct():
    start = time.process_time()
    hs300 = get_data()
    # print(hs300.head(5))
    print(type(hs300))
    df = ZIG(hs300, 0.1, 'close')[::]
    df.query('pv_is!=0 or pv_ok!=0').to_csv('zig.csv')
    peak = df.query('peak_is==1')['peak']
    valley = df.query('valley_is==1')['valley']
    pv_ok = df.query('pv_ok!=0')
    diff_p = peak.diff(1)
    diff_v = valley.diff(1)
    pv = df.query('valley_is==1 or peak_is==1')['pv']
    pv_ok_value = df.query('pv_ok!=0')['pv_ok_value']
    df2 = ZIG(hs300, 0.1 * 2.5, 'close')[::]
    pv2 = df2.query('valley_is==1 or peak_is==1')['pv']
    end = time.process_time()
    print("%.2f s" % (end - start))

    print(len(pv), len(pv2), len(pv_ok))
    print(pv_ok_value)
    # print(diff_p, diff_v)

    # df = ZIG(hs300, 0.1, 'close')
    ama = AMA(hs300['close'])
    amah = AMA(hs300['high'])
    amal = AMA(hs300['low'])
    # print(df[df['peak_is']==1 or df['valley_is']==1])
    # valley = df[df['valley_is'] == 1]
    # peak = df[df['peak_is'] == 1]
    # # pv = df[(df['valley_is'==1]) | (df['peak_is']==1)]
    # # print(pv)
    # pv = df.query('valley_is==1 or peak_is==1')
    # # print(df.query('valley_is==1 or peak_is==1'))
    # # print(peak.loc[:,['high','low','close']])
    # # print(valley.loc[:,['high','low','close']])
    #
    # sumpeak = peak['close'].sum()
    #
    # sump1 = sumpeak - peak['close'][0]
    # print('sumpeak=', sumpeak)
    # sumvalley = valley['close'].sum()
    # sumv0 = sumvalley - valley['close'][-1]
    # sumv1 = sumvalley - valley['close'][0]
    #
    # pvl = [len(peak), len(valley)]
    #
    # pmv = sumpeak - sumvalley
    # vmp = sump1 - sumv0
    # # print(pvl, pmv, vmp)
    # # print(peak.loc[:, ['high', 'low', 'close']])
    # # print(valley.loc[:, ['high', 'low', 'close']])
    #
    # btr = cal_back_test_result(df, sgn_datetime='date',
    #                            sgn_buy='valley_is', sgn_sell='peak_is',
    #                            sgn_short='valley_is', sgn_cover='peak_is',
    #                            open_long_cond='valley_is==1', close_long_cond='peak_is==1',
    #                            open_short_cond='peak_is==1', close_short_cond='valley_is==1',
    #                            sgn_vol='volume', multiplier=1, slip_points=0)
    # # btr['nv']=np.exp(btr['nv'])
    # print(btr['longTotal'], btr['tradeTotal'], btr['tradeWinRatio'], np.max(btr['nv']))
    # # sumRP =  [(p-v)/v for p,v in [peak['close'],valley['close']]].sum()
    # sumRP = sum([v for v in ((peak['close'][i + 1] - valley['close'][i]) / valley['close'][i] for i in
    #                          range(len(peak['close']) - 1))])
    # sumRV = sum(
    #     [v for v in
    #      ((peak['close'][i] - valley['close'][i]) / peak['close'][i] for i in range(len(peak['close']) - 1))])
    # profitRP = [2 - pv['close'][i + 1] / pv['close'][i] for i in range(len(pv['close']) - 1) if pv['peak_is'][i] == 1]
    # profitRV = [pv['close'][i + 1] / pv['close'][i] for i in range(len(pv['close']) - 1) if pv['valley_is'][i] == 1]
    #
    # print('ZIG sumRp=', sumRP, 'sumRV=', sumRV)
    # print('ZIG profitRP=', np.cumprod(profitRP))
    # print('ZIG profitRv=', np.cumprod(profitRV))
    # # peak['NP']=copy.copy(peak['close'])
    # # for i in range(len(peak)-1):
    # #     peak.loc[peak.index[i],'NP']=profitRV[i]
    # tc = TrendClass(df, timeWindow=60, rangePCT=20000, riskTimes=2.5,
    #                 dif_times=4.0, sgnDatetime='date', sgnHigh='high', sgnLow='low')  # ,sgnHigh='close',sgnLow='close'
    # fr = Fractal(df, periods=9)
    # tcList = [item.__dict__ for item in tc]
    # tcdf = pd.DataFrame(tcList).set_index('datetime')
    # tcdf['close'] = df.close
    # # btr1 = CalBackTestResult(tcdf, sgnDatetime='date',
    # #                          sgn_buy='isUp', sgn_sell='isDown',
    # #                          sgn_short='isDown', sgn_cover='isUp',
    # #                          openLongCond='isUp==1', closeLongCond='isDown==1',
    # #                          openShortCond='isDown==1', closeShortCond='isUp==1',
    # #                          sgn_vol='close', multiplier=1, slip_points=0)
    # btr1 = cal_back_test_result(df, sgn_datetime='date', sgn_trade_price='open',
    #                             sgn_buy='lok', sgn_sell='uok',
    #                             sgn_short='uok', sgn_cover='lok',
    #                             open_long_cond='lok==1', close_long_cond='uok==1',
    #                             open_short_cond='uok==1', close_short_cond='lok==1',
    #                             sgn_vol='open', multiplier=1, slip_points=0)
    # # print('trend strategy result:', btr1)
    # aa = tcdf[tcdf['isUp'] == 1]  # [(td.cprice,td.top,td.bot) for td in tc]
    # bb = tcdf[tcdf['isDown'] == 1]  # [(td.top,td.bot) for td in tc]
    # cc = [td.cprice for td in tc]
    #
    # ulOK = tcdf.query('upOk + downOk >= 1')
    # # print('ulOK',len(ulOK),ulOK[['close','top','bot', 'isUp', 'isDown', 'isBalance','upOk','downOk' ]])
    # ul = tcdf.query('isUp + isDown == 1')
    # ulb = tcdf.query('isUp +isDown + isBalance==1')
    # # print(ulb[['close', 'isUp', 'isDown', 'isBalance']])
    # dfu = tcdf.query('isUp == 1 ')
    # dfl = tcdf.query(' isDown == 1')
    # profitShort = [2 - ulb['close'][i + 1] / ulb['close'][i]
    #                for i in range(len(ulb['close']) - 1) if ulb['isDown'][i] == 1]
    # profitLong = [ulb['close'][i + 1] / ulb['close'][i]
    #               for i in range(len(ulb['close']) - 1) if ulb['isUp'][i] == 1]
    #
    # sumPS = np.cumprod(profitShort)
    # sumPL = np.cumprod(profitLong)
    # print('TRENAD profitRVp=', np.cumprod([item for item in profitLong if item > 1]))
    # meanU = dfu.close.mean()
    # meanL = dfl.close.mean()
    #
    # # print(ulb[['close', 'uprice', 'lprice', 'isUp', 'isDown', 'isBalance']].head(5))
    # # print(ulb[['close', 'uprice', 'lprice', 'isUp', 'isDown', 'isBalance']].tail(5))
    # print(len(ulb), meanU, len(dfu), '\n', meanL, len(dfl), len(ul),
    #       '\nprofit_short=', sumPS, '\nprofit_long=', sumPL)
    # ul1 = [(td.uprice1, td.lprice1) for td in tc]
    #
    # difu = tcdf.query('difu>0 and isUp==1')
    # # print(difu.tail(5),difu.head(5))
    #
    # zig = df
    # peak = zig['peak']
    # valley = zig['valley']
    # peak_is = zig['peak_is']
    # valley_is = zig['valley_is']
    # tcdf = df
    # print('longCount=', btr1['longCount'], 'shortCount=', btr1['shortCount'],
    #       'nv_var=', np.sqrt(250) * np.std(np.log(tcdf['nv']).diff(1)),
    #       '\nDATA_STD=', np.sqrt(250) * np.std(np.log(tcdf['close']).diff(1)),
    #       '\nnv=', tcdf['nv'][-1], '\nlong_nv=', tcdf['long_nv'][-1],
    #       '\nshort_nv=', tcdf['short_nv'][-1],
    #       '\nlongDrawback=', max(btr1['longDrawback']), max(btr1['longDrawback1']),
    #       '\nshortDrawback=', max(btr1['shortDrawback']), max(btr1['shortDrawback1']),
    #       '\ndrawback=', max(btr1['drawback1']),
    #       '\nlong_nv_list=', tcdf.query('uok==1')['long_nv']
    #       )

    # print(btr1['long_trade'][['close','nv','long_nv','short_nv']],
    #       '\n',btr1['short_trade'][['close','nv','long_nv','short_nv']],
    #       '\n',tcdf['short_nv'].tail(10),tcdf['long_nv'].tail(10))
    # isUp = tcdf['difu']
    # isDown = tcdf['difl']
    # difs = isUp + isDown  # print(len(aa),len(bb),len(ul),'\n')

    # print(df.head(100)[['peak_is','valley_is','close','nv']])
    import matplotlib.pyplot as plt
    import matplotlib.path as mpath

    # pc = plt.subplot(3, 1, 1)

    # pc.plot(peak)
    # pc.plot(valley)
    plt.plot(df['close'])

    # plt.plot(pv2, marker=mpath.Path.unit_circle())
    plt.plot(pv, marker=mpath.Path.unit_regular_star(2))
    plt.plot(pv_ok_value, marker=mpath.Path.unit_circle())
    plt.legend()
    plt.show()
    # pc.plot(valley)
    # pc.plot(isUp)
    # pc.plot(isDown)
    # pc1 = plt.subplot(3, 1, 2)
    # pc1.plot(pv['close'])
    # pc1.plot(ul)
    # pc1.plot(ul['cprice'])
    # pc1.plot(ul['uprice'])
    # pc1.plot(ulb['close'])
    # pc1.plot(tcdf['cprice'])
    # pc2 = plt.subplot(3, 1, 3)
    # pc2.plot(difs)
    # pc.plot(close)
    # pc1.plot(df['close'])

    # ps = plt.subplot(3, 1, 1)
    # ps.plot(ulb['close'])
    # ps.plot(pv['close'])
    # ps.plot(np.exp(np.log(df['close']).diff(1).cumsum()))
    #
    # ps1 = plt.subplot(3,1,2)
    # ps1.plot(tcdf['nv'])
    # ps1.plot(np.exp(np.log(df['close']).diff(1).cumsum()))

    # win = pg.GraphicsWindow()
    #
    # p = win.addPlot()
    # ps1 = p.plot(pen='r')
    # ps2 = plt.subplot(2, 1, 2)
    # ps2.plot(tcdf['nv'])
    # ps2.plot(tcdf['long_nv'])
    # ps2.plot(tcdf['short_nv'][-500:-1])
    # hs300=hs300.reset_index()  #([i for i in range(len(hs300.index))])
    # hs300['ama'] = ama
    # hs300['amah'] = amah
    # hs300['amal'] = amal
    # ps1.setData(hs300.index, hs300['close'])
    # win.show()
    print(df.head(5))
    # print(hs300.head(5))
    p = 66
    # ps1.plot(hs300['amah'].rolling(window=p,min_periods=1).max(),color='r')
    # ps1.plot(hs300['amal'].rolling(window=p, min_periods=1).min(), color='b')

    # ps1.plot(pv)
    # # ps1.plot(df2['peak'])
    # ps1.plot(pv2)

    # plt.legend()
    # #
    # plt.show()


if __name__ == '__main__':
    # from ctpquant.dbhelper.sqlitedb import SqliteHelper
    ct()

    """
    from matplotlib.finance import  candlestick2_ohlc
    import matplotlib.pyplot as plt
    from pymongo import MongoClient
    from pymongo.errors import ConnectionFailure
    from ui_tools.plotK import *

    from datetime import datetime
    # dbfile="./data/quotedb.db"
    # db=SqliteHelper(dbfile)
    # sql="select InstrumentID as symbol,TradeDay as date,OpenPrice as open,
    #     HighPrice as high ,LowPrice as low ,CLosePrice as close ,
    #     volume,Holding as openInterest,TradeDateTime as datetime
    #      from basedata{0} where InstrumentID='{1}' and tradedatetime>'{2}'
    #      order by  TradeDateTime asc".format('m1','RB1705','2017-01-11 00:00:00')
    # df=db.ExecuteSqlFrame(sql)
    port = 27017
    host = 'localhost'
    try:
        dbClient = MongoClient(host, port, serverSelectionTimeoutMS=500)
        dbClient.server_info()
        # print(dbClient)
        dbName = 'CTP_Tick_Db'
        collectionName = 'ru'
        cond = {'date': {'$gte': '20170818'},'$and':[{'symbol':'ru1801'}]}
        db = dbClient[dbName]
        collection = db[collectionName]
        output = { 'open': 1, 'high': 1, 'low': 1, 'close': 1, 'volume': 1,'datetime':1}  #'date': 1, 'time': 1,
        output = {'lastPrice':1,'datetime':1}
        cursor = collection.find(cond, output)
        dftick = pd.DataFrame()
        # if cursor:
        dftick = pd.DataFrame(list(cursor))
        del dftick['_id']
        df0=dftick.set_index('datetime')
        # print(type(df0['datetime'][0]))

        # print(df, type(df['date']))
        # data=[df.datetime,df.open,df.high,df.low,df.close,df.openInterest]
        #
        # high=df['high'].values
        # close=df['close'].values
        # low=df['low'].values
        # dt=df.datetime.values
        # fg=plot_k(df,volFieldName='volume')
        # fg.show()
        # print(df0.head(5))
        dfLast = df0.resample('1Min').ohlc()
        df = dfLast.lastPrice.dropna(axis=0,how='any')
        # del df[[df.open]]

        # print(df.index[4])
        tc=TrendClass(df, timeWindow=30, rangePCT=35000, riskTimes=1.5)
        tcdf = pd.DataFrame(tc)
        # print(tcdf)
        aa = [(td.cprice,td.top,td.bot) for td in tc]
        bb =  [(td.top,td.bot) for td in tc]
        cc = [td.cprice for td in tc]
        ul = [(td.uprice,td.lprice) for td in tc]
        ul1 = [(td.uprice1, td.lprice1) for td in tc]

        close = df['close']
        # print(type(list(close)),type(close))
        high = df['high']
        low = df['low']
        open = df['open']
        # vol = df['volume']
        dt=df.index
        # print(dt[0],len(low))
        frc=Fractal(dt,high,low,2)
        ufp=[frc[i].uprice1 for i in range(len(frc))]
        lfp = [frc[i].lprice1 for i in range(len(frc))]
        spH=Speed(dt,high,5)

        zig =ZIG(df,0.05,'lh')
        peak = zig['peak']
        valley = zig['valley']
        peak_is = zig['peak_is']
        valley_is = zig['valley_is']
        # print(close)
        tr,atr = TR(df)
        # print(matr) 
        # tr=matr['TR']
        # atr = matr['ATR']

        pc = plt.subplot(3, 1, 1)


        pc.plot(peak)
        pc.plot(valley)
        pc1 = plt.subplot(3, 1, 2)
        # pc1.plot(ul)
        pc1.plot(peak_is)
        pc1.plot(valley_is)
        pc2=plt.subplot(3,1,3)
        pc2.plot(cc)
        # pc.plot(close)

        # pc.plot(zw.HHV(high, n=15))  # ['hhh'])
        # pc.plot(zw.LLV(low, n=15))  # ['lll'])

        # pc.plot(tr)
        # pc.plot(atr)

        # pc1=plt.subplot(3,1,2)
        # # pc1.plot(MA((close),20))
        # # pc1.plot(MA((close), 100))
        # pc2 = plt.subplot(3, 1, 3)
        # pc2.plot(ufp)
        # pc2.plot(lfp)
        # pc2.plot(close)
        # pc2.plot(zw.EMA(close,20))

        # candlestick2_ohlc(pc1, df.open,df.high,df.low,df.close, width=0.6,colorup='r', colordown='g')
        candlestick2_ohlc(pc2, df.open, df.high, df.low, df.close, width=0.6, colorup='r', colordown='g')

        plt.show()
    except ConnectionFailure as e:
        print(e)
    """
